from typing import TypedDict, Annotated
from langgraph.graph import StateGraph, END
from nodes import NovelNodes, NovelState
from util import writeFile

class NovelWorkflow:
    def __init__(self):
        self.nodes = NovelNodes()
        self.graph = StateGraph(NovelState)
        
    def _should_refine_outline(self, state: NovelState):
        return state['refinement_count'] < 5

    def compile_workflow(self):
        # 添加工作流节点
        self.graph.add_node("generate_outline", self.nodes.generate_initial_outline)
        self.graph.add_node("refine_outline", self.nodes.refine_outline)
        self.graph.add_node("write_chapter", lambda state: self.nodes.write_chapter(state, state["current_chapter"]))
        self.graph.add_node("final_review", self.nodes.final_review)

        # 设置初始流程
        self.graph.set_entry_point("generate_outline")
        
        # 条件分支
        self.graph.add_conditional_edges(
            "generate_outline",
            lambda state: "refine_outline" if self._should_refine_outline(state) else "write_chapter"
        )
        
        # 连接节点
        self.graph.add_edge("refine_outline", "write_chapter")
        
        self.graph.add_conditional_edges(
            "write_chapter",
            lambda state: "write_chapter" if state.get("current_chapter") else "final_review"
        )
        
        self.graph.add_edge("final_review", END)
        
        # 配置递归深度
        return self.graph.compile()